##独立的RSI计算器，直接使用xtdata获取当日历史分钟数据
from xtquant import xtdata
from typing import Dict, Optional, Union
import pandas as pd
from datetime import datetime
import time


class IndependentRSICalculator:
    """
    完全独立的RSI计算器
    
    特性：
    1. 不依赖订阅，直接使用xtdata获取历史数据
    2. 输入股票代码即可获取最新RSI（上一分钟结束时的RSI）
    3. 自动下载当日分钟数据并计算
    4. 支持缓存，避免重复下载
    
    使用方法：
        rsi_calc = IndependentRSICalculator(period=5)
        rsi = rsi_calc.get_latest_rsi("513050.SH")  # 直接获取最新RSI
    """
    
    def __init__(self, period: int = 5, xtdata=None):
        if period <= 0:
            raise ValueError("RSI period must be positive")
        self.period: int = period
        self.xtdata = xtdata
        
        # 每个股票的状态
        self._last_close: Dict[str, float] = {}
        self._sma_up: Dict[str, float] = {}
        self._sma_down: Dict[str, float] = {}
        self._last_rsi: Dict[str, float] = {}
        
        # 缓存机制：避免重复下载同一天的数据
        self._cache_date: Dict[str, str] = {}  # 股票代码 -> 缓存日期
        self._cache_ttl: int = 60  # 缓存有效期（秒）- 1分钟更新一次确保获取最新完整分钟数据
        self._cache_time: Dict[str, float] = {}  # 股票代码 -> 缓存时间戳
        self._cache_minute: Dict[str, str] = {}  # 股票代码 -> 缓存的分钟时间戳
        
        # 交易时间段常量
        self.MORNING_START = datetime.strptime("09:30", "%H:%M").time()
        self.MORNING_END = datetime.strptime("11:30", "%H:%M").time()
        self.AFTERNOON_START = datetime.strptime("13:00", "%H:%M").time()
        self.AFTERNOON_END = datetime.strptime("15:00", "%H:%M").time()
    
    def reset(self, stock_code: Optional[str] = None) -> None:
        """重置指定股票或全部股票的RSI状态"""
        if stock_code is None:
            self._last_close.clear()
            self._sma_up.clear()
            self._sma_down.clear()
            self._last_rsi.clear()
            self._cache_date.clear()
            self._cache_time.clear()
            self._cache_minute.clear()
            return
        
        self._last_close.pop(stock_code, None)
        self._sma_up.pop(stock_code, None)
        self._sma_down.pop(stock_code, None)
        self._last_rsi.pop(stock_code, None)
        self._cache_date.pop(stock_code, None)
        self._cache_time.pop(stock_code, None)
        self._cache_minute.pop(stock_code, None)
    
    def _is_cache_valid(self, stock_code: str, today: str) -> bool:
        """检查缓存是否有效"""
        if stock_code not in self._cache_date or stock_code not in self._cache_time:
            return False
        
        # 检查日期是否匹配
        if self._cache_date[stock_code] != today:
            return False
        
        # 检查当前分钟是否与缓存分钟相同
        current_minute = datetime.now().strftime("%Y%m%d%H%M")
        cached_minute = self._cache_minute.get(stock_code, "")
        if cached_minute != current_minute:
            return False
        
        # 检查时间是否在有效期内
        current_time = time.time()
        cache_time = self._cache_time[stock_code]
        return (current_time - cache_time) < self._cache_ttl
    
    def _download_minute_data(self, stock_code: str, date: str) -> bool:
        """下载指定股票的当日分钟数据"""
        if not self.xtdata:
            print("⚠️ xtdata未初始化，无法下载数据")
            return False
        
        try:
            self.xtdata.download_history_data(
                stock_code=stock_code, 
                period="1m", 
                start_time=date, 
                end_time=date
            )
            time.sleep(0.1)  # 避免请求过快
            return True
        except Exception as e:
            print(f"⚠️ {stock_code} 分钟数据下载失败: {e}")
            return False
    
    def _get_minute_data(self, stock_code: str, date: str) -> Optional[pd.DataFrame]:
        """获取指定股票的当日分钟数据"""
        if not self.xtdata:
            return None
        
        try:
            data = self.xtdata.get_local_data(
                stock_list=[stock_code], 
                period='1m', 
                start_time=date, 
                end_time=date
            )
            
            df = data.get(stock_code) if data else None
            if df is None or df.empty:
                return None
            
            # 识别时间列和价格列
            time_col = next((c for c in df.columns if 'time' in c.lower()), None)
            price_col = next((c for c in ['close', 'lastPrice', 'price', 'last', 'current'] if c in df.columns), None)
            
            if not time_col or not price_col:
                return None
            
            # 时间格式统一处理
            if str(df[time_col].dtype) in ('int64', 'float64'):
                df['__time__'] = pd.to_datetime(df[time_col], unit='ms') + pd.Timedelta(hours=8)
            else:
                try:
                    df['__time__'] = pd.to_datetime(df[time_col], format='%Y%m%d%H%M%S', errors='coerce')
                except Exception:
                    df['__time__'] = pd.to_datetime(df[time_col], errors='coerce')
            
            df = df[df['__time__'].notna()].sort_values('__time__')
            
            # 添加列名到DataFrame以便后续使用
            df._price_col = price_col
            
            return df if not df.empty else None
            
        except Exception as e:
            print(f"⚠️ {stock_code} 分钟数据获取失败: {e}")
            return None
    
    def _filter_trading_session_data(self, df: pd.DataFrame) -> pd.DataFrame:
        """筛选交易时段内的数据，并排除当前未完成的分钟"""
        try:
            # 先筛选交易时段内的数据
            mask = [((self.MORNING_START <= ts.time() <= self.MORNING_END) or
                     (self.AFTERNOON_START <= ts.time() <= self.AFTERNOON_END)) for ts in df['__time__']]
            if any(mask):
                filtered_df = df[mask]
                
                # 排除当前未完成的分钟：只保留完整的分钟数据
                current_time = datetime.now()
                current_minute = current_time.replace(second=0, microsecond=0)
                
                # 只保留时间小于当前分钟的数据（即已完成的分钟）
                complete_mask = filtered_df['__time__'] < current_minute
                if complete_mask.any():
                    return filtered_df[complete_mask]
                else:
                    return filtered_df  # 如果没有完整分钟数据，返回原数据
            else:
                return df
        except Exception:
            pass
        return df
    
    def _calculate_rsi_from_data(self, stock_code: str, df: pd.DataFrame) -> Optional[float]:
        """从分钟数据计算RSI"""
        try:
            # 安全地获取价格列名
            price_col = getattr(df, '_price_col', None)
            if not price_col:
                # 如果没有_price_col属性，重新识别价格列
                price_col = next((c for c in ['close', 'lastPrice', 'price', 'last', 'current'] if c in df.columns), None)
                if not price_col:
                    print(f"⚠️ {stock_code} 无法识别价格列")
                    return None
            
            prices = df[price_col].dropna().tolist()
            
            if len(prices) < 2:
                return None
            
            # 重置该股票的状态
            self._last_close[stock_code] = prices[0]
            self._sma_up[stock_code] = 0.0
            self._sma_down[stock_code] = 0.0
            
            # 依次处理每个价格
            for price in prices[1:]:
                self._update_rsi_state(stock_code, price)
            
            return self._last_rsi.get(stock_code)
            
        except Exception as e:
            print(f"⚠️ {stock_code} RSI计算异常: {e}")
            return None
    
    def _update_rsi_state(self, stock_code: str, close_price: float) -> None:
        """更新RSI内部状态（标准RSI公式）"""
        close_price = float(close_price)
        
        if stock_code not in self._last_close:
            self._last_close[stock_code] = close_price
            return
        
        change = close_price - self._last_close[stock_code]
        up = change if change > 0 else 0.0
        down = abs(change) if change < 0 else 0.0
        
        if stock_code not in self._sma_up or stock_code not in self._sma_down:
            self._sma_up[stock_code] = up
            self._sma_down[stock_code] = down
        else:
            n = self.period
            self._sma_up[stock_code] = (self._sma_up[stock_code] * (n - 1) + up) / n
            self._sma_down[stock_code] = (self._sma_down[stock_code] * (n - 1) + down) / n
        
        self._last_close[stock_code] = close_price
        
        # 标准RSI公式：RSI = 100 - (100 / (1 + RS))
        sma_down = self._sma_down[stock_code]
        if sma_down == 0:
            rsi = 100.0 if self._sma_up[stock_code] > 0 else 50.0
        else:
            rs = self._sma_up[stock_code] / sma_down
            rsi = 100.0 - (100.0 / (1.0 + rs))
        
        self._last_rsi[stock_code] = rsi
    
    def _filter_data_before_time(self, df: pd.DataFrame, end_time: datetime) -> pd.DataFrame:
        """
        过滤掉指定时间之后的数据（回测模式使用）
        
        Args:
            df: 原始数据DataFrame
            end_time: 结束时间
            
        Returns:
            过滤后的DataFrame
        """
        try:
            if df.empty:
                return df
            
            # 确保时间列存在
            if '__time__' not in df.columns:
                return df
            
            # 过滤掉结束时间之后的数据
            filtered_df = df[df['__time__'] <= end_time]
            return filtered_df
            
        except Exception as e:
            print(f"⚠️ 数据过滤异常: {e}")
            return df
    
    def get_latest_rsi(self, stock_code: str, end_time: Optional[datetime] = None) -> Optional[float]:
        """
        获取指定股票的最新RSI值（上一分钟结束时的RSI）
        
        Args:
            stock_code: 股票代码
            end_time: 结束时间（回测模式使用，实时模式为None）
            
        Returns:
            RSI值（0-100），如果无法计算则返回None
        """
        try:
            # 确定当前日期
            if end_time is not None:
                # 回测模式：使用指定的结束时间
                today = end_time.strftime("%Y%m%d")
                current_time = end_time
            else:
                # 实时模式：使用当前时间
                today = datetime.now().strftime("%Y%m%d")
                current_time = datetime.now()
            
            # 检查缓存是否有效
            if self._is_cache_valid(stock_code, today):
                return self._last_rsi.get(stock_code)
            
            # 缓存无效，需要重新下载和计算
            print(f"🔄 {stock_code} 开始下载当日分钟数据...")
            
            # 下载分钟数据
            if not self._download_minute_data(stock_code, today):
                return None
            
            # 获取分钟数据
            df = self._get_minute_data(stock_code, today)
            if df is None:
                return None
            
            # 筛选交易时段内的数据
            df = self._filter_trading_session_data(df)
            if df.empty:
                return None
            
            # 在回测模式下，过滤掉结束时间之后的数据
            if end_time is not None:
                df = self._filter_data_before_time(df, end_time)
                if df.empty:
                    return None
            
            # 计算RSI
            rsi = self._calculate_rsi_from_data(stock_code, df)
            
            # 更新缓存
            if rsi is not None:
                self._cache_date[stock_code] = today
                self._cache_time[stock_code] = time.time()
                self._cache_minute[stock_code] = current_time.strftime("%Y%m%d%H%M")
                ##print(f"✅ {stock_code} RSI计算完成: {rsi:.2f}")
            else:
                ##print(f"⚠️ {stock_code} RSI计算失败")
                pass
            
            return rsi
            
        except Exception as e:
            print(f"❌ {stock_code} 获取RSI异常: {e}")
            return None
    
    def get_rsi_with_preview(self, stock_code: str, current_price: Optional[float] = None) -> Dict[str, Optional[float]]:
        """
        获取RSI值，包括确认值和预览值
        
        Args:
            stock_code: 股票代码
            current_price: 当前价格（用于预览RSI）
            
        Returns:
            包含确认RSI和预览RSI的字典
        """
        # 获取确认RSI
        confirmed_rsi = self.get_latest_rsi(stock_code)
        
        # 计算预览RSI
        preview_rsi = None
        if current_price is not None and confirmed_rsi is not None:
            preview_rsi = self._calculate_preview_rsi(stock_code, current_price)
        
        return {
            'confirmed': confirmed_rsi,
            'preview': preview_rsi
        }
    
    def _calculate_preview_rsi(self, stock_code: str, current_price: float) -> Optional[float]:
        """计算预览RSI（基于当前价格）"""
        try:
            if stock_code not in self._last_close:
                return None
            
            n = self.period
            change = float(current_price) - float(self._last_close[stock_code])
            up = change if change > 0 else 0.0
            down = abs(change) if change < 0 else 0.0
            
            if stock_code not in self._sma_up or stock_code not in self._sma_down:
                sma_up_prev = 0.0
                sma_down_prev = 0.0
            else:
                sma_up_prev = float(self._sma_up[stock_code])
                sma_down_prev = float(self._sma_down[stock_code])
            
            sma_up_tmp = (sma_up_prev * (n - 1) + up) / n
            sma_down_tmp = (sma_down_prev * (n - 1) + down) / n
            
            if sma_down_tmp == 0:
                return 100.0 if sma_up_tmp > 0 else 50.0
            
            rs = sma_up_tmp / sma_down_tmp
            return 100.0 - (100.0 / (1.0 + rs))
            
        except Exception:
            return None
    
    def force_refresh(self, stock_code: str) -> Optional[float]:
        """强制刷新指定股票的RSI数据"""
        self.reset(stock_code)
        return self.get_latest_rsi(stock_code)
    
    def get_all_stock_rsi(self, stock_codes: list) -> Dict[str, Optional[float]]:
        """批量获取多个股票的RSI值"""
        results = {}
        for code in stock_codes:
            results[code] = self.get_latest_rsi(code)
        return results


# 为了保持向后兼容，保留原有的类名
StandardRSICalculator = IndependentRSICalculator

# 更新导出列表
__all__ = ["IndependentRSICalculator", "StandardRSICalculator"]


# 使用示例
if __name__ == "__main__":
    # 定义股票代码
    stock_code = "513120.SH"  # 可以修改为您需要的股票代码
    # 创建RSI计算器（周期为5）
    rsi_calc = IndependentRSICalculator(period=5,xtdata=xtdata)
    # 获取最新RSI
    target_time = datetime(2025, 8, 22, 14, 30, 1)  # 14:30:00
    rsi_value = rsi_calc.get_latest_rsi(stock_code,target_time)
    if rsi_value is not None:
        print(f"股票 {stock_code} 的最新RSI: {rsi_value:.2f}")
    else:
        print(f"无法获取 {stock_code} 的RSI值")
